package org.zlui65.milan.processor

import org.zlui65.milan.annotations.Style
import org.zlui65.milan.annotations.Styleable
import org.zlui65.milan.annotations.StyleableChild
import java.io.File
import javax.annotation.processing.AbstractProcessor
import javax.annotation.processing.ProcessingEnvironment
import javax.annotation.processing.RoundEnvironment
import javax.annotation.processing.SupportedAnnotationTypes
import javax.annotation.processing.SupportedSourceVersion
import javax.lang.model.SourceVersion
import javax.lang.model.element.Element
import javax.lang.model.element.ElementKind
import javax.lang.model.element.TypeElement
import javax.tools.Diagnostic

/**
 * Annotation processor for generating style extension functions.
 * This processor handles @Styleable, @StyleableChild, and @Style annotations.
 */
@SupportedAnnotationTypes(
    "org.zlui65.milan.annotations.Styleable",
    "org.zlui65.milan.annotations.StyleableChild",
    "org.zlui65.milan.annotations.Style"
)
@SupportedSourceVersion(SourceVersion.RELEASE_17)
class StyleProcessor : AbstractProcessor() {

    override fun process(annotations: Set<TypeElement>, roundEnv: RoundEnvironment): Boolean {
        if (annotations.isEmpty()) return false

        // Process @Styleable classes
        val styleableElements = roundEnv.getElementsAnnotatedWith(Styleable::class.java)
        for (element in styleableElements) {
            if (element.kind != ElementKind.CLASS) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.ERROR,
                    "@Styleable can only be applied to classes",
                    element
                )
                continue
            }

            processStyleableClass(element as TypeElement)
        }

        // Process @Style fields
        val styleElements = roundEnv.getElementsAnnotatedWith(Style::class.java)
        for (element in styleElements) {
            if (element.kind != ElementKind.FIELD) {
                processingEnv.messager.printMessage(
                    Diagnostic.Kind.ERROR,
                    "@Style can only be applied to fields",
                    element
                )
                continue
            }

            processStyleField(element)
        }

        return true
    }

    private fun processStyleableClass(element: TypeElement) {
        val packageName = processingEnv.elementUtils.getPackageOf(element).toString()
        val className = element.simpleName.toString()
        val styleableName = element.getAnnotation(Styleable::class.java).value

        // Generate extension functions for this styleable class
        generateStyleableExtensions(packageName, className, styleableName, element)
    }

    private fun processStyleField(element: Element) {
        val enclosingElement = element.enclosingElement as TypeElement
        val packageName = processingEnv.elementUtils.getPackageOf(enclosingElement).toString()
        val className = enclosingElement.simpleName.toString()
        val fieldName = element.simpleName.toString()

        // Generate extension functions for this style field
        generateStyleExtensions(packageName, className, fieldName, element)
    }

    private fun generateStyleableExtensions(
        packageName: String,
        className: String,
        styleableName: String,
        element: TypeElement
    ) {
        val fileName = "${className}StyleExtensions"
        val fileContent = buildString {
            append("package $packageName\n\n")
            append("import org.zlui65.milan.Style\n")
            append("import org.zlui65.milan.StyleBuilder\n")
            append("import org.zlui65.milan.style\n\n")
            
            // Find all @StyleableChild fields
            val styleableChildFields = element.enclosedElements.filter { 
                it.kind == ElementKind.FIELD && it.getAnnotation(StyleableChild::class.java) != null 
            }
            
            for (field in styleableChildFields) {
                val fieldName = field.simpleName.toString()
                val attrId = field.getAnnotation(StyleableChild::class.java).value
                
                // Generate extension function for setting this child's style
                append("/**\n")
                append(" * Sets the style for the $fieldName child view.\n")
                append(" */\n")
                append("fun StyleBuilder.${fieldName}Style(styleRes: Int) {\n")
                append("    setAttribute($attrId, styleRes)\n")
                append("}\n\n")
                
                // Generate extension function for building this child's style
                append("/**\n")
                append(" * Builds a style for the $fieldName child view.\n")
                append(" */\n")
                append("fun StyleBuilder.${fieldName}Style(builder: StyleBuilder.() -> Unit) {\n")
                append("    val styleBuilder = StyleBuilder()\n")
                append("    styleBuilder.builder()\n")
                append("    setAttribute($attrId, styleBuilder.build())\n")
                append("}\n\n")
            }
        }
        
        // Write the generated file
        writeToFile(packageName, fileName, fileContent)
    }

    private fun generateStyleExtensions(
        packageName: String,
        className: String,
        fieldName: String,
        element: Element
    ) {
        val fileName = "${className}StyleExtensions"
        val fileContent = buildString {
            append("package $packageName\n\n")
            append("import org.zlui65.milan.Style\n")
            append("import org.zlui65.milan.StyleBuilder\n")
            append("import org.zlui65.milan.style\n\n")
            
            // Generate extension function for adding this style
            append("/**\n")
            append(" * Adds the $fieldName style to the builder.\n")
            append(" */\n")
            append("fun StyleBuilder.add${fieldName.capitalize()}() {\n")
            append("    add($className.$fieldName)\n")
            append("}\n\n")
        }
        
        // Write the generated file
        writeToFile(packageName, fileName, fileContent)
    }

    private fun writeToFile(packageName: String, fileName: String, content: String) {
        val kaptKotlinGeneratedDir = processingEnv.options["kapt.kotlin.generated"]
            ?: return

        val packageDir = File(kaptKotlinGeneratedDir).resolve(
            packageName.replace('.', File.separatorChar)
        )
        packageDir.mkdirs()

        val file = packageDir.resolve("$fileName.kt")
        file.writeText(content)
    }

    private fun String.capitalize(): String {
        return if (this.isEmpty()) this
        else this[0].uppercaseChar() + this.substring(1)
    }
}
